from typing import Dict, List, Optional, Union

import numpy as np
from torch import Tensor

from .summarizer import Summarizer


class MAPSummarizer(Summarizer):
    def __init__(self, class_names: List[str]):
        """
        :param class_names:List[str]
            类别名称，如车辆检测的["car", "bus", "truck"]
        """
        super().__init__()
        self.class_names = class_names

    def _init_metric_cache(self):
        self.gt_cls_cache = []
        self.metric_cache = []  # 列表中的每个元素为由correct_flag, conf, cls构成的ndarray

    def calculate_metrics(self) -> Optional[Dict[str, Union[Tensor, float]]]:
        metric_cache = np.concatenate(self.metric_cache)
        metric_cache = metric_cache[np.argsort(-metric_cache[:, 1])]
        gt_cls_cache = np.concatenate(self.gt_cls_cache)
        metrics = dict()
        ap_sum = 0
        for cls in np.unique(gt_cls_cache):
            ap = self.calculate_ap(
                total_objects=(gt_cls_cache == cls).sum(),
                metric_cache=metric_cache[metric_cache[:, 2] == cls, 0],
            )
            metrics[self.class_names[int(cls)]] = ap
            ap_sum += ap
        metrics["mAP"] = ap_sum / len(metrics)
        return metrics

    def update_metric_cache(self, result, label):
        self.gt_cls_cache.append(self.get_cls_cache(label))
        if len(result) * len(label) == 0:
            return
        frame_cache = self.get_frame_cache(result, label)
        if frame_cache:
            self.metric_cache.append(np.array(frame_cache))

    def get_frame_cache(self, result, label_cpy):
        raise NotImplementedError

    @staticmethod
    def get_cls_cache(label):
        raise NotImplementedError

    @staticmethod
    def calculate_ap(total_objects, metric_cache):
        if total_objects == 0:
            return -1
        # get pr lane
        numerator = np.cumsum(metric_cache)
        p = numerator / np.arange(1, len(metric_cache) + 1)
        r = numerator / total_objects
        # calculate area
        p0, r0 = 1, 0
        double_s = 0
        for p1, r1 in zip(p, r):
            double_s += (p0 + p1) * (r1 - r0)
            p0, r0 = p1, r1
        return double_s / 2
